# frozen_string_literal: true

require 'spec_helper'

RSpec.describe Gitlab::Llm::Completions::ExplainVulnerability, feature_category: :vulnerability_management do
  let(:prompt_class) { Gitlab::Llm::Templates::ExplainVulnerability }
  let(:example_answer) { "Sure, ..." }
  let(:example_response) do
    {
      Gitlab::Llm::VertexAi::Client => {
        "predictions" => [
          {
            "candidates" => [
              {
                "author" => "",
                "content" => example_answer
              }
            ],
            "safetyAttributes" => {
              "categories" => ["Violent"],
              "scores" => [0.4000000059604645],
              "blocked" => false
            }
          }
        ],
        "deployedModelId" => "1",
        "model" => "projects/1/locations/us-central1/models/codechat-bison-001",
        "modelDisplayName" => "codechat-bison-001",
        "modelVersionId" => "1"
      }.to_json,
      Gitlab::Llm::Anthropic::Client => example_answer
    }
  end

  let(:error_response) do
    {
      Gitlab::Llm::VertexAi::Client => { error: { message: 'Ooops...' } }.to_json,
      Gitlab::Llm::Anthropic::Client => nil
    }
  end

  let(:errors) do
    {
      Gitlab::Llm::VertexAi::Client => ["Ooops..."],
      Gitlab::Llm::Anthropic::Client => ["The response from the AI provider was empty."]
    }
  end

  let_it_be(:user) { create(:user) }
  let_it_be(:project) do
    create(:project, :custom_repo, files: {
      'main.c' => "#include <stdio.h>\n\nint main() { printf(\"hello, world!\"); }"
    })
  end

  let_it_be(:vulnerability) { create(:vulnerability, :with_finding, project: project) }

  subject(:explain) { described_class.new(prompt_class, { request_id: 'uuid' }) }

  before do
    allow(GraphqlTriggers).to receive(:ai_completion_response)
    vulnerability.finding.location['file'] = 'main.c'
    vulnerability.finding.location['start_line'] = 1
  end

  context 'when a null prompt is returned by the template class' do
    before do
      allow_next_instance_of(prompt_class) do |prompt_class|
        allow(prompt_class).to receive(:to_prompt).and_return(nil)
      end
    end

    it 'returns the default error response' do
      explain.execute(user, vulnerability, {})

      expect(GraphqlTriggers).to have_received(:ai_completion_response)
        .with({ user_id: user.to_global_id, resource_id: vulnerability.to_global_id }, hash_including({
          id: anything,
          model_name: vulnerability.class.name,
          content: '',
          role: ::Gitlab::Llm::Cache::ROLE_ASSISTANT,
          request_id: 'uuid',
          errors: [described_class::NULL_PROMPT_ERROR]
        }))
    end
  end

  shared_examples 'explain vulnerability completions' do |llm_client, client_method|
    context 'when the client returns an unsuccessful response' do
      before do
        allow_next_instance_of(llm_client) do |client|
          allow(client).to receive(client_method).and_return(
            error_response[llm_client]
          )
        end
      end

      it 'publishes the error to the graphql subscription' do
        explain.execute(user, vulnerability, {})

        expect(GraphqlTriggers).to have_received(:ai_completion_response)
          .with({ user_id: user.to_global_id, resource_id: vulnerability.to_global_id }, hash_including({
            id: anything,
            model_name: vulnerability.class.name,
            content: '',
            role: ::Gitlab::Llm::Cache::ROLE_ASSISTANT,
            request_id: 'uuid',
            errors: errors[llm_client]
          }))
      end
    end

    context 'when the client returns a successful response' do
      before do
        allow(llm_client).to receive(:new).and_call_original
        allow_next_instance_of(llm_client) do |client|
          allow(client).to receive(client_method).and_return(example_response[llm_client])
        end
      end

      it 'sets the vertex client to retry content blocked requests' do
        explain.execute(user, vulnerability, {})

        expect(llm_client).to have_received(:new)
          .with(user, **extra_arguments)
      end

      it 'publishes the content from the AI response' do
        explain.execute(user, vulnerability, {})

        expect(GraphqlTriggers).to have_received(:ai_completion_response)
          .with({ user_id: user.to_global_id, resource_id: vulnerability.to_global_id }, hash_including({
            id: anything,
            model_name: vulnerability.class.name,
            content: example_answer,
            role: ::Gitlab::Llm::Cache::ROLE_ASSISTANT,
            request_id: 'uuid',
            errors: []
          }))
      end

      context 'when an unexpected error is raised' do
        let(:error) { StandardError.new("Ooops...") }

        before do
          allow_next_instance_of(llm_client) do |client|
            allow(client).to receive(client_method).and_raise(error)
          end
          allow(Gitlab::ErrorTracking).to receive(:track_exception)
        end

        it 'records the error' do
          explain.execute(user, vulnerability, {})
          expect(Gitlab::ErrorTracking).to have_received(:track_exception).with(error)
        end

        it 'publishes a generic error to the graphql subscription' do
          explain.execute(user, vulnerability, {})

          expect(GraphqlTriggers).to have_received(:ai_completion_response)
            .with({ user_id: user.to_global_id, resource_id: vulnerability.to_global_id }, hash_including({
              id: anything,
              model_name: vulnerability.class.name,
              content: '',
              role: ::Gitlab::Llm::Cache::ROLE_ASSISTANT,
              request_id: 'uuid',
              errors: [described_class::DEFAULT_ERROR]
            }))
        end
      end

      context 'when the client experiences a Net::ReadTimeout' do
        let(:error) { Net::ReadTimeout.new }

        before do
          allow_next_instance_of(llm_client) do |client|
            allow(client).to receive(client_method).and_raise(error)
          end
          allow(Gitlab::ErrorTracking).to receive(:track_exception)
        end

        it 'records the error' do
          explain.execute(user, vulnerability, {})
          expect(Gitlab::ErrorTracking).to have_received(:track_exception).with(error)
        end

        it 'publishes a generic error to the graphql subscription' do
          explain.execute(user, vulnerability, {})

          expect(GraphqlTriggers).to have_received(:ai_completion_response)
            .with({ user_id: user.to_global_id, resource_id: vulnerability.to_global_id }, hash_including({
              id: anything,
              model_name: vulnerability.class.name,
              content: '',
              role: ::Gitlab::Llm::Cache::ROLE_ASSISTANT,
              request_id: 'uuid',
              errors: [described_class::CLIENT_TIMEOUT_ERROR]
            }))
        end
      end

      context 'when request is cached', :use_clean_rails_redis_caching do
        context 'when `include_source_code` is not used' do
          before do
            allow_next_instance_of(llm_client) do |client|
              allow(client).to receive(client_method).once.and_return(example_response[llm_client])
            end

            explain.execute(user, vulnerability, {})
          end

          it 'executes the request just once' do
            expect(llm_client).not_to receive(:new)

            explain.execute(user, vulnerability, {})

            expect(GraphqlTriggers).to have_received(:ai_completion_response)
              .with({ user_id: user.to_global_id, resource_id: vulnerability.to_global_id }, hash_including({
                id: anything,
                model_name: vulnerability.class.name,
                content: example_answer,
                role: ::Gitlab::Llm::Cache::ROLE_ASSISTANT,
                request_id: 'uuid',
                errors: []
              })).twice
          end
        end

        context 'when the `include_source_code` option is toggled' do
          before do
            allow_next_instance_of(llm_client) do |client|
              allow(client).to receive(client_method).twice.and_return(example_response[llm_client])
            end
          end

          it 'bypasses the cache and make a fresh request' do
            # cache miss
            explain.execute(user, vulnerability, { include_source_code: true })
            explain.execute(user, vulnerability, { include_source_code: false })
            explain.execute(user, vulnerability, {})

            # cache hit
            explain.execute(user, vulnerability, { include_source_code: true })
            explain.execute(user, vulnerability, { include_source_code: false })
            explain.execute(user, vulnerability, {})

            expect(GraphqlTriggers).to have_received(:ai_completion_response)
              .with(
                { user_id: user.to_global_id, resource_id: vulnerability.to_global_id },
                hash_including({
                  id: anything,
                  model_name: vulnerability.class.name,
                  content: example_answer,
                  role: ::Gitlab::Llm::Cache::ROLE_ASSISTANT,
                  request_id: 'uuid',
                  errors: []
                })
              ).exactly(6).times
          end
        end
      end
    end
  end

  describe '#execute', :clean_gitlab_redis_cache do
    let(:tracking_context) { { request_id: "uuid" } }

    context 'when the explain_vulnerability_anthropic feature flag is enabled' do
      it_behaves_like 'explain vulnerability completions', Gitlab::Llm::Anthropic::Client, :stream do
        let(:extra_arguments) { { tracking_context: tracking_context } }
      end
    end

    context 'when the explain_vulnerability_anthropic feature flag is disabled' do
      before do
        stub_feature_flags(explain_vulnerability_anthropic: false)
      end

      it_behaves_like 'explain vulnerability completions', Gitlab::Llm::VertexAi::Client, :text do
        let(:extra_arguments) { { retry_content_blocked_requests: true, tracking_context: tracking_context } }
      end
    end
  end
end
